import pandas as pd
import model.config as conf


def get_data(path):
    data_list = []
    if path.endswith('.csv'):
        df = pd.read_csv(path, dtype='object')
        for index, row in df.iterrows():
            # label = row.label
            # if label not in conf.cls.keys():
            #     label = '1-红色'
            #     data_list.append((row.context, int(conf.cls[label])))
            # else:
            #     data_list.append((row.context, int(conf.cls[row.label])))
            data_list.append((row.context, int(conf.cls[row.label])))
    elif path.endswith('.txt'):
        with open(path, encoding='utf-8') as f:
            for line in f.readlines():
                result = line.strip().split(' ')
                data_list.append((result[0],  int(result[1])))
    elif path.endswith('.xlsx'):
        df = pd.read_excel(path, engine='openpyxl', dtype='object')
        for index, row in df.iterrows():
            label = row.label
            context = '。'.join([context[3:] for context in row.context.split('\n')])
            # if label not in conf.cls.keys():
            #     label = '1-红色'
            data_list.append((row.context, int(conf.cls[label])))
    else:
        raise Exception('请使用.txt或csv格式文件作为输入，txt格式为 “content + 空格 + label”， csv格式为 “content,label”')

    return df, data_list
